from nd_utils.venue_re_util import *

class CoorgCluster():
    def __init__(self, paperidx_venue_dict, paperidx_list):
        self.paperidx_venue_dict = paperidx_venue_dict
        self.paperidx_list = paperidx_list

        self.cluster_idx_venue_dict = self._get_cluster_idx_venue_dict()

        self.cluster_venue_idx_dict = self._get_cluster_venue_idx_dict()

        self.cluster_venue_set = set(self.cluster_venue_idx_dict.keys())

    def _get_cluster_idx_venue_dict(self):
        cluster_idx_venue_dict = {}
        for idx in self.paperidx_list:
            venue = self.paperidx_venue_dict[idx]
            cluster_idx_venue_dict[idx] = get_clean_paper_venue(venue)
        return cluster_idx_venue_dict

    def _get_cluster_venue_idx_dict(self):
        cluster_venue_idx_dict = {}
        for idx, venue in self.cluster_idx_venue_dict.items():
            if venue != None:
                if venue not in cluster_venue_idx_dict:
                    cluster_venue_idx_dict[venue] = [ idx ]
                else:
                    cluster_venue_idx_dict[venue].append(idx)
        return cluster_venue_idx_dict

    def get_cluster_idx_by_venue(self, venue):
        return self.cluster_venue_idx_dict[venue][0]